import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision
import torch.optim as optim
import torchvision.models as models
import onnx
import torch.onnx as torch_onnx
from mmdet.apis import init_detector

config_file = '../../configs/anti-uav/ssd300_voc.py'
checkpoint_file = '../work_dirs/ssd300_voc/latest.pth'

# model = models.resnet50(pretrained=True)
model = init_detector(config_file, checkpoint_file)

# Load the weights from a file (.pth usually)
# weights_path = './mnist.pth'
# state_dict = torch.load(weights_path)
#
# # Load the weights now into a model net architecture defined by our class
# model.load_state_dict(state_dict)

# Create the right input shape (e.g. for an image)
input = torch.randn(1, 3, 224, 224)

torch.onnx.export(model, input, "ssd300_voc.onnx", verbose=True)



# Load the ONNX model
model = onnx.load("ssd300_voc.onnx")

# Check that the IR is well formed
onnx.checker.check_model(model)

# Print a human readable representation of the graph
onnx.helper.printable_graph(model.graph)
